import torch
import torch.nn.functional as F
from torchvision.utils import save_image
import os
from tqdm import tqdm
import matplotlib.pyplot as plt
import numpy as np
import math

from model.model import SPIN
from data import get_dataloaders
from train import Config
cfg = Config()

# 模型路径


# 创建输出目录
os.makedirs("results/concat", exist_ok=True)

# 初始化模型
model = SPIN(
    height=cfg.height,
    width=cfg.width,
    grid_size=cfg.grid_size,
    num_iters=cfg.num_iters,
    hidden_dim=cfg.hidden_dim,
    num_heads=cfg.num_heads,
    patch_size=cfg.patch_size_model,
    overlap=cfg.overlap,
    num_blocks=cfg.num_blocks,
    upscale_factor=cfg.scale_factor
).to(cfg.device)

# 加载参数
model_path = "checkpoints/best_model.pth"
input_size=torch.rand([1, 3, cfg.height, cfg.width], device=cfg.device)
model(input_size)   # 走一遍，确保参数加载
checkpoint = torch.load(model_path, map_location=cfg.device)
model.load_state_dict(checkpoint)
print(f"✅ 模型权重已加载: {model_path}")

# 数据
_, test_loader = get_dataloaders(
    batch_size=1,
    num_workers=cfg.num_workers
)

# 将tensor转为numpy
def tensor_to_np(img):
    img = img.squeeze(0).cpu().clamp(0, 1).numpy()
    img = np.transpose(img, (1, 2, 0))  # CxHxW → HxWxC
    return img

# 计算 PSNR (dB) 指标
def psnr(target, prediction, max_val=1.0):
    mse = torch.mean((target - prediction) ** 2)
    psnr_value = 20 * torch.log10(max_val / torch.sqrt(mse))
    return psnr_value.item()

# 推理
model.eval()
total=0
with torch.no_grad():
    for idx, batch in enumerate(tqdm(test_loader)):
        lr = batch['lr'].to(cfg.device)
        hr = batch['hr'].to(cfg.device)

        # 模型输出
        sr = model(lr)

        # 双线性插值（Bilinear Interpolation）
        lr_bilinear = F.interpolate(lr, size=hr.shape[-2:], mode='bilinear', align_corners=False)

        # 计算 PSNR
        psnr_sr = psnr(hr, sr)
        total+=psnr_sr
        psnr_bilinear = psnr(hr, lr_bilinear)

        # 转换为dB（可以选择输出dB值）
        print(f"PSNR (SR Model) = {psnr_sr:.2f} dB")
        print(f"PSNR (Bilinear) = {psnr_bilinear:.2f} dB")

        # 将图片转换为 numpy
        lr_img = tensor_to_np(lr)
        hr_img = tensor_to_np(hr)
        sr_img = tensor_to_np(sr)
        bilinear_img = tensor_to_np(lr_bilinear)

        # 拼接可视化（不resize）
        fig, axs = plt.subplots(1, 4, figsize=(20, 5))
        axs[0].imshow(hr_img)
        axs[0].set_title(f"HR (GT)\nPSNR={psnr_sr:.2f} dB")
        axs[0].axis("off")

        axs[1].imshow(lr_img)
        axs[1].set_title("LR (Input)")
        axs[1].axis("off")

        axs[2].imshow(sr_img)
        axs[2].set_title(f"SR (Model)\nPSNR={psnr_sr:.2f} dB")
        axs[2].axis("off")

        axs[3].imshow(bilinear_img)
        axs[3].set_title(f"Bilinear\nPSNR={psnr_bilinear:.2f} dB")
        axs[3].axis("off")

        plt.tight_layout()
        plt.savefig(f"results/concat/{idx:04d}.png")
        plt.close()

        if idx >= 19:
            break
print(total/len(test_loader))
print("✅ 拼接图像和PSNR保存完成（不做上采样）")
